package top.sqdpt.rag.service.impl;


import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;
import top.sqdpt.rag.assistant.RagAssistant;
import top.sqdpt.rag.service.AiChatService;

import java.util.List;

@Service
public class AiChatServiceImpl implements AiChatService {
    @Autowired
    private RagAssistant ragAssistant;

    @Autowired
    private EmbeddingStore<TextSegment> embeddingStore;

    @Autowired
    private EmbeddingModel embeddingModel;

    @Override
    public String chat(String chatId, String userMsg) {
        return ragAssistant.chat(chatId, userMsg);
    }

    @Override
    public void loadDoc(String filePath) {
        Document document = FileSystemDocumentLoader.loadDocument(filePath, new ApacheTikaDocumentParser());
        if (document == null) {
            return;
        }
        DocumentSplitter splitter = DocumentSplitters.recursive(800, 200);
        List<TextSegment> segments = splitter.split(document);
        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
        embeddingStore.addAll(embeddings);
    }


    @Override
    public void loadDocByDir(String dirPath) {
        List<Document> documents = FileSystemDocumentLoader.loadDocuments(dirPath, new ApacheTikaDocumentParser());
        if (CollectionUtils.isEmpty(documents)) {
            return;
        }
        DocumentSplitter splitter = DocumentSplitters.recursive(800, 200);
        List<TextSegment> segments = splitter.splitAll(documents);
        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
        embeddingStore.addAll(embeddings);
    }

}
